import io
import json
import flask
import torch
import torch.nn.functional as F
from PIL import Image
from torchvision import transforms as T
from torchvision.models import resnet50
from torch.autograd import Variable

app = flask.Flask(__name__)
model = None
use_gpu = False

with open("imagenet_class.txt") as f:
    idx2label = eval(f.read())

def load_model():
    global model
    model = resnet50(pretrained=True)
    model.eval()
    if use_gpu:
        model.cuda()

def prepare_image(image, target_size):
    if image.mode != "RGB":
        image = image.convert("RGB")

    image = T.Resize(target_size)(image)
    image = T.ToTensor()(image)

    image = T.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])(image)

    image = image[None]
    if use_gpu:
        image = image.cuda()
    return Variable(image, volatile=True)

# 开启服务端
@app.route("/predict", methods=["POST"])
def predict():
    data = {"success": False}

    if flask.request.method == "POST":
        if flask.request.files.get("image"):
            image = flask.request.files["image"].read()
            image = Image.open(io.BytesIO(image))

            image = prepare_image(image,target_size=(224,224))

            preds = F.softmax(model(image),dim=1)
            results = torch.topk(preds.cpu().data,k=3,dim=1)
            results = (results[0].cpu().numpy(),results[1].cpu().numpy())

            data["predictions"] = list()

            for prob,label in zip(results[0][0],results[1][0]):
                label_name = idx2label[label]
                r = {"label":label_name,"probability":float(prob)}
                data["predictions"].append(r)

            data["success"] = True
    return flask.jsonify(data)



if __name__ == "__main__":
    load_model()
    app.run()
